#!/usr/bin/python
import cv
import sys
from random import uniform
from time import sleep
from math import sin,cos,pi,atan2,sqrt
from numpy import matrix
from time import time
from random import randrange


capture = cv.CreateCameraCapture( 0 )
cv.NamedWindow( "Fig", 1 )
frame = cv.QueryFrame( capture )
SizeX,SizeY=cv.GetSize(frame)
den=2
segment= cv.CreateImage((SizeX/den,SizeY/den), 8, 3 )
segment2= cv.CreateImage((SizeX/den,SizeY/den), 8, 3 )
segment3= cv.CreateImage((SizeX/den,SizeY/den), 8, 3 )
segment4= cv.CreateImage((SizeX/den,SizeY/den), 8, 3 )
segment5= cv.CreateImage((SizeX/den,SizeY/den), 8, 3 )
segmentc= cv.CreateImage((SizeX/den,SizeY/den), 8, 3 )
hsv= cv.CreateImage((SizeX/den,SizeY/den), 8, 3 )
dst= cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
dst2= cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
d= cv.CreateImage((SizeX/den,SizeY/den), cv.IPL_DEPTH_16S, 1 )
d2=cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
d3=cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
blackLines= cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
both= cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
harr= cv.CreateImage((SizeX/den,SizeY/den), 32, 1 )
W,H=SizeX/den,SizeY/den
lastdetected= 0
THR=100
dects=50

hue= cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
sat= cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )
val= cv.CreateImage((SizeX/den,SizeY/den), 8, 1 )


prevface=[(0,0),(5,0),(0,5)]
dodetection=True
def grid(p):
    return int(round(p[0])), int(round(p[1]))
    
def ptdst(point1,point2):
    return sqrt((point1[0]-point2[0])*(point1[0]-point2[0])+(point1[1]-point2[1])*(point1[1]-point2[1]))
    
def ptdstw(point1,point2):
    if point1[1]<100 or point2[1]<100:
        return 300.0+abs(point1[0]-point2[0])
    else:
        return abs(point1[0]-point2[0])
        
    
def ptdst3(point1,point2):
    dist=sqrt((point1[0]-point2[0])*(point1[0]-point2[0])+(point1[1]-point2[1])*(point1[1]-point2[1])+(point1[2]-point2[2])*(point1[2]-point2[2]))
    if (point1[0]>245 and point1[1]>245 and point1[2]>245):
        dist=dist+300.0
    return dist
    

def avg(point1,point2):
    return (0.5*point1[0]+0.5*point2[0], 0.5*point2[1]+0.5*point2[1])

def near(t1,t2,t):
    return abs(t1[0]-t2[0])<t and abs(t1[1]-t2[1])<t

def winded(point1,point2,p3,p4):
    avg=(0.25*(point1[0]+point2[0]+p3[0]+p4[0]),0.25*(point1[1]+point2[1]+p3[1]+p4[1]))
    ps=[(atan2(p[1]-avg[1], p[0]-avg[0]), p) for p in [point1,point2,p3,p4]]
    ps.sort(reverse=True)
    return [p[1] for p in ps]
 
def neighbors(f,s):
    if f==0 and s==0: return ((1,2),(4,0))
    if f==0 and s==1: return ((4,3),)
    if f==0 and s==2: return ((4,6),(3,0))
    if f==0 and s==3: return ((1,5),)
    if f==0 and s==5: return ((3,3),)
    if f==0 and s==6: return ((1,8),(5,2))
    if f==0 and s==7: return ((5,5),)
    if f==0 and s==8: return ((3,6),(5,8))
    
    if f==1 and s==0: return ((2,2),(4,2))
    if f==1 and s==1: return ((4,1),)
    if f==1 and s==2: return ((4,0),(0,0))
    if f==1 and s==3: return ((2,5),)
    if f==1 and s==5: return ((0,3),)
    if f==1 and s==6: return ((2,8),(5,0))
    if f==1 and s==7: return ((5,1),)
    if f==1 and s==8: return ((0,6),(5,2))
    
    if f==2 and s==0: return ((4,8),(3,2))
    if f==2 and s==1: return ((4,5),)
    if f==2 and s==2: return ((4,2),(1,0))
    if f==2 and s==3: return ((3,5),)
    if f==2 and s==5: return ((1,3),)
    if f==2 and s==6: return ((3,8),(5,6))
    if f==2 and s==7: return ((5,3),)
    if f==2 and s==8: return ((1,6),(5,0))
    
    if f==3 and s==0: return ((4,6),(0,2))
    if f==3 and s==1: return ((4,7),)
    if f==3 and s==2: return ((4,8),(2,0))
    if f==3 and s==3: return ((0,5),)
    if f==3 and s==5: return ((2,3),)
    if f==3 and s==6: return ((0,8),(5,8))
    if f==3 and s==7: return ((5,7),)
    if f==3 and s==8: return ((2,6),(5,6))
    
    if f==4 and s==0: return ((1,2),(0,0))
    if f==4 and s==1: return ((1,1),)
    if f==4 and s==2: return ((1,0),(2,2))
    if f==4 and s==3: return ((0,1),)
    if f==4 and s==5: return ((2,1),)
    if f==4 and s==6: return ((0,2),(3,0))
    if f==4 and s==7: return ((3,1),)
    if f==4 and s==8: return ((3,2),(2,0))
    
    if f==5 and s==0: return ((1,6),(2,8))
    if f==5 and s==1: return ((1,7),)
    if f==5 and s==2: return ((1,8),(0,6))
    if f==5 and s==3: return ((2,7),)
    if f==5 and s==5: return ((0,7),)
    if f==5 and s==6: return ((2,6),(3,8))
    if f==5 and s==7: return ((3,7),)
    if f==5 and s==8: return ((3,6),(0,8))
    
def processColors(useRGB=True):
    global assigned,didassignments

    bestj=0
    besti=0
    bestcon=0
    match_position=0
    bestd=10001
    taken=[0 for i in range(6)]
    done=0
    opposite={0:2, 1:3, 2:0, 3:1, 4:5, 5:4}
    poss={}
    for j,f in enumerate(hsvs):
        for i,s in enumerate(f):
            poss[j,i]=range(6)
    

    toloop=hsvs
    if useRGB: toloop=colors
    
    while done<48:
        bestd=10000
        forced=False
        for j,f in enumerate(toloop):
            for i,s in enumerate(f):
                if i!=4 and assigned[j][i]==-1 and (not forced):
                    considered=0
                    for k in poss[j,i]:
                        if taken[k]<8:
                            
                            if useRGB:
                                dist=ptdst3(s, toloop[k][4])
                            else:
                                dist=ptdstw(s, toloop[k][4])
                                
                            considered+=1
                            if dist<bestd:
                                bestd=dist
                                bestj=j
                                besti=i
                                match_position=k
                    if besti==i and bestj==j: bestcon=considered
                    if considered==1:
                        forced=True
        done=done+1
        assigned[bestj][besti]=match_position
        print bestcon
        
        op= opposite[match_position]
        ns=neighbors(bestj,besti)
        for neighbor in ns:
            p=poss[neighbor]
            if match_position in p: p.remove(match_position)
            if op in p: p.remove(op)
        
        taken[match_position]+=1
        
    didassignments=True
    
def compfaces(f1,f2):
    ansd=0
    for point1 in f1:
        mind=1000000000
        for point2 in f2:
            d=ptdst(point1,point2)
            if d<mind:
                mind=d
        ansd += mind
    return ansd/4

succ=0
tracking=0
win_size=5
flags=0
detected=0

grey = cv.CreateImage ((W,H), 8, 1)
prev_grey = cv.CreateImage ((W,H), 8, 1)
pyramid = cv.CreateImage ((W,H), 8, 1)
prev_pyramid = cv.CreateImage ((W,H), 8, 1)

ff= cv.InitFont(cv.CV_FONT_HERSHEY_PLAIN, 1,1, shear=0, thickness=1, lineType=8)

counter=0
undetectednum=100
stage=1
extract=False
selected=0
colors=[[] for i in range(6)]
hsvs=[[] for i in range(6)]
assigned=[[-1 for i in range(9)] for j in range(6)]
for i in range(6):
    assigned[i][4]=i
    
didassignments=False

mycols=[(0,127,255), (20,240,20), (0,0,255), (200,0,0), (0,255,255), (255,255,255)]
def intersect_seg(x1,x2,x3,x4,y1,y2,y3,y4):
    den= (y4-y3)*(x2-x1)-(x4-x3)*(y2-y1)
    if abs(den)<0.1: return (False, (0,0),(0,0))
    ua=(x4-x3)*(y1-y3)-(y4-y3)*(x1-x3)
    ub=(x2-x1)*(y1-y3)-(y2-y1)*(x1-x3)    
    ua=ua/den
    ub=ub/den
    x=x1+ua*(x2-x1)
    y=y1+ua*(y2-y1)
    return (True,(ua,ub),(x,y))

while True:
    frame = cv.QueryFrame( capture )
    if not frame:
        cv.WaitKey(0)
        break
    
    cv.Resize( frame, segment)
    
    cv.Copy(segment, segmentc)

    cv.CvtColor (segment, grey, cv.CV_RGB2GRAY)
    
    if tracking>0:
        
        detected=2
        features, status, track_error = cv.CalcOpticalFlowPyrLK (
            prev_grey, grey, prev_pyramid, pyramid,
            features,
            (win_size, win_size), 3,
            (cv.CV_TERMCRIT_ITER|cv.CV_TERMCRIT_EPS, 20, 0.03),
            flags)
        features = [ p for (st,p) in zip(status, features) if st]
        if len(features)<4: 
            tracking= 0
        else:
            ds1=ptdst(features[0], features[1])
            ds2=ptdst(features[2], features[3])
            if max(ds1,ds2)/min(ds1,ds2)>1.4: tracking=0
            
            ds3=ptdst(features[0], features[2])
            ds4=ptdst(features[1], features[3])
            if max(ds3,ds4)/min(ds3,ds4)>1.4: tracking=0
            
            if ds1< 10 or ds2<10 or ds3<10 or ds4<10: tracking=0
            if tracking==0: detected=0
            
    if tracking==0:
        detected=0
        cv.Smooth(grey,dst2,cv.CV_GAUSSIAN, 3)
        cv.Laplace(dst2,d)
        cv.CmpS(d,8,d2,cv.CV_CMP_GT)
        if lastdetected>dects: THR=THR+1
        if lastdetected<dects: THR=max(2,THR-1)
        li= cv.HoughLines2(d2, cv.CreateMemStorage(), cv.CV_HOUGH_PROBABILISTIC, 1, 3.1415926/45, THR, 10, 5)
        angs=[]
        for (point1, point2) in li: 
            a = atan2(point2[1]-point1[1],point2[0]-point1[0])
            if a<0:a+=pi
            angs.append(a)
            
        t=10
        totry=[]
        for i in range(len(li)):
            point1,point2=li[i]
            for j in range(i+1,len(li)):
                q1,q2=li[j]
                
                dd1= sqrt((point2[0]-point1[0])*(point2[0]-point1[0])+(point2[1]-point1[1])*(point2[1]-point1[1]))
                dd2= sqrt((q2[0]-q1[0])*(q2[0]-q1[0])+(q2[1]-q1[1])*(q2[1]-q1[1]))
                if max(dd1,dd2)/min(dd1,dd2)>1.3: continue
                
                matched=0
                if near(point1,q2,t): 
                    IT=(avg(point1,q2), point2, q1,dd1)
                    matched=matched+1
                if near(point2,q2,t): 
                    IT=(avg(point2,q2), point1, q1,dd1)
                    matched=matched+1
                if near(point1,q1,t): 
                    IT=(avg(point1,q1), point2, q2,dd1)
                    matched=matched+1
                if near(point2,q1,t): 
                    IT=(avg(point2,q1), q2, point1,dd1)
                    matched=matched+1
                
                if matched==0:
                    point1=(float(point1[0]),float(point1[1]))
                    point2=(float(point2[0]),float(point2[1]))
                    q1=(float(q1[0]),float(q1[1]))
                    q2=(float(q2[0]),float(q2[1]))
                    success,(ua,ub),(x,y)=\
                        intersect_seg(point1[0],point2[0],q1[0],q2[0],point1[1],point2[1],q1[1],q2[1])
                    if success and ua>0 and ua<1 and ub>0 and ub<1:
                        ok1=0
                        ok2=0
                        if abs(ua-1.0/3)<0.05:ok1=1
                        if abs(ua-2.0/3)<0.05:ok1=2
                        if abs(ub-1.0/3)<0.05:ok2=1
                        if abs(ub-2.0/3)<0.05:ok2=2
                        if ok1>0 and ok2>0:
                            if ok1==2: point1,point2=point2,point1
                            if ok2==2: q1,q2=q2,q1
                            z1=(q1[0]+2.0/3*(point2[0]-point1[0]),q1[1]+2.0/3*(point2[1]-point1[1]))
                            z2=(point1[0]+2.0/3*(q2[0]-q1[0]),point1[1]+2.0/3*(q2[1]-q1[1]))
                            z=(point1[0]-1.0/3*(q2[0]-q1[0]),point1[1]-1.0/3*(q2[1]-q1[1]))
                            IT=(z,z1,z2,dd1)
                            matched=1
                if matched==1:
                    a1 = atan2(point2[1]-point1[1],point2[0]-point1[0])
                    a2 = atan2(q2[1]-q1[1],q2[0]-q1[0])
                    if a1<0:a1+=pi
                    if a2<0:a2+=pi
                    ang=abs(abs(a2-a1)-pi/2)
                    if ang < 0.5:
                        totry.append(IT)
               
        
        res=[]
        for i in range(len(totry)):
            
            p,point1,point2,dd=totry[i]
            a1 = atan2(point1[1]-p[1],point1[0]-p[0])
            a2 = atan2(point2[1]-p[1],point2[0]-p[0])
            if a1<0:a1+=pi
            if a2<0:a2+=pi
            dd=1.7*dd
            evidence=0
            totallines=0
            
            A = matrix([[point2[0]-p[0],point1[0]-p[0],p[0]],[point2[1]-p[1],point1[1]-p[1],p[1]],[0,0,1]])
            Ainv= A.I
            
            v=matrix([[point1[0]],[point1[1]],[1]])
            
            for j in range(len(li)):
            
                a = angs[j]
                ang1=abs(abs(a-a1)-pi/2)
                ang2=abs(abs(a-a2)-pi/2)
                if ang1 > 0.1 and ang2>0.1: continue
                
                q1,q2= li[j]
                qwe=0.06
                
                v=matrix([[q1[0]],[q1[1]],[1]])
                vp=Ainv*v;
                if vp[0,0] > 1.1 or vp[0,0]<-0.1: continue
                if vp[1,0] > 1.1 or vp[1,0]<-0.1: continue
                if abs(vp[0,0]-1/3.0)>qwe and abs(vp[0,0]-2/3.0)>qwe and \
                    abs(vp[1,0]-1/3.0)>qwe and abs(vp[1,0]-2/3.0)>qwe: continue
                
                v=matrix([[q2[0]],[q2[1]],[1]])
                vp=Ainv*v;
                if vp[0,0] > 1.1 or vp[0,0]<-0.1: continue
                if vp[1,0] > 1.1 or vp[1,0]<-0.1: continue
                if abs(vp[0,0]-1/3.0)>qwe and abs(vp[0,0]-2/3.0)>qwe and \
                    abs(vp[1,0]-1/3.0)>qwe and abs(vp[1,0]-2/3.0)>qwe: continue
                    
                evidence+=1
            
            res.append((evidence, (p,point1,point2)))
        
        minch=10000
        res.sort(reverse=True)
        if len(res)>0:
            minps=[]
            point=[]
            for i in range(len(res)):
                if res[i][0]>0.05*dects:
                    p,point1,point2=res[i][1]
                    p3= (point2[0]+point1[0]-p[0], point2[1]+point1[1]-p[1])
                    w=[p,point1,point2,p3]
                    p3= (prevface[2][0]+prevface[1][0]-prevface[0][0], 
                         prevface[2][1]+prevface[1][1]-prevface[0][1])
                    tc= (prevface[0],prevface[1],prevface[2],p3)
                    ch=compfaces(w,tc) 
                    if ch<minch:
                        minch=ch
                        minps= (p,point1,point2)

            if len(minps)>0:
                prevface=minps
                
                if minch<10:
                    succ+=1
                    point=prevface
                    detected=1
                    
            else:
                succ=0
            
            if succ>2 and 1:
                point=[]
                for i in [1.0/3, 2.0/3]:
                    for j in [1.0/3, 2.0/3]:
                        point.append((p0[0]+i*Dist_first[0]+j*Dist_second[0], p0[1]+i*Dist_first[1]+j*Dist_second[1]))
                
                features=point
                tracking=1
                succ=0
                
    else:
        p=features[0]
        point1=features[1]
        point2=features[2]
        
        Dist_first=(point1[0]-p[0],point1[1]-p[1])
        Dist_second=(point2[0]-p[0],point2[1]-p[1])
        
        point=[(p[0]-Dist_first[0]-Dist_second[0], p[1]-Dist_first[1]-Dist_second[1]),
            (p[0]+2*Dist_second[0]-Dist_first[0], p[1]+2*Dist_second[1]-Dist_first[1]),
            (p[0]+2*Dist_first[0]-Dist_second[0], p[1]+2*Dist_first[1]-Dist_second[1])]
        
        prevface=[point[0],point[1],point[2]]

    if (detected or undetectednum<1) and dodetection:
        if not detected: 
            undetectednum+=1
            point=lastpt
        if detected: 
            undetectednum=0
            lastpt=point
        
        cv.CvtColor(segmentc, hsv, cv.CV_RGB2HSV)
        cv.Split(hsv, hue, sat, val, None)
        p3= (point[2][0]+point[1][0]-point[0][0], point[2][1]+point[1][1]-point[0][1])
        cv.Line(segment,grid(point[0]),grid(point[1]),(0,255,0),2)
        cv.Line(segment,grid(point[1]),grid(p3),(0,255,0),2)
        cv.Line(segment,grid(p3),grid(point[2]),(0,255,0),2)
        cv.Line(segment,grid(point[2]),grid(point[0]),(0,255,0),2)
        point=winded(point[0],point[1],point[2],p3)
        Dist_first=(point[1][0]-point[0][0], point[1][1]-point[0][1])
        Dist_second=(point[3][0]-point[0][0], point[3][1]-point[0][1])
        p0=(point[0][0],point[0][1])
        
        ep=[]
        midpts=[]
        i=1
        j=5
        for k in range(9):
            ep.append((p0[0]+i*Dist_first[0]/6.0+j*Dist_second[0]/6.0, p0[1]+i*Dist_first[1]/6.0+j*Dist_second[1]/6.0))
            i=i+2
            if i==7:
                i=1
                j=j-2
        
        rad= ptdst(Dist_first,(0.0,0.0))/6.0
        cs=[]
        hsvcs=[]
        den=2
        
        for i,p in enumerate(ep):
            if p[0]>rad and p[0]<W-rad and p[1]>rad and p[1]<H-rad:
                color=cv.Avg(segmentc[int(p[1]-rad/den):int(p[1]+rad/den),int(p[0]-rad/den):int(p[0]+rad/den)])
                color=cv.Avg(segmentc[int(p[1]-rad/den):int(p[1]+rad/den),int(p[0]-rad/den):int(p[0]+rad/den)])
                cv.Circle(segment, grid(p), int(rad), color,-1)
                if i==4:
                    cv.Circle(segment, grid(p), int(rad), (0,255,255),2)
                else:
                    cv.Circle(segment, grid(p), int(rad), (255,255,255),2)
                
                hueavg= cv.Avg(hue[int(p[1]-rad/den):int(p[1]+rad/den),int(p[0]-rad/den):int(p[0]+rad/den)])
                satavg= cv.Avg(sat[int(p[1]-rad/den):int(p[1]+rad/den),int(p[0]-rad/den):int(p[0]+rad/den)])
                if extract:
                    cs.append(color)
                    hsvcs.append((hueavg[0],satavg[0]))
                    
        if extract:
            extract= not extract
            colors[selected]=cs
            hsvs[selected]=hsvcs
            selected=min(selected+1,5)
    x=20
    y=20
    s=13
    for i in range(6):
        if not colors[i]: 
            x+=3*s+10
            continue
        cv.Rectangle(segment, (x-1,y-1), (x+3*s+5,y+3*s+5), (0,0,0),-1)
        x1,y1=x,y
        x2,y2=x1+s,y1+s
        for j in range(9):
            if didassignments:
                cv.Rectangle(segment, (x1,y1), (x2,y2), mycols[assigned[i][j]],-1)
            else:
                cv.Rectangle(segment, (x1,y1), (x2,y2), colors[i][j],-1)
            x1+=s+2
            if j==2 or j==5: 
                x1=x
                y1+=s+2
            x2,y2=x1+s,y1+s
        x+=3*s+10
    
    x=20
    y=20
    for i in range(selected):x+=3*s+10
    cv.Rectangle(segment, (x-1,y-1), (x+3*s+5,y+3*s+5), (0,0,255),2)

    lastdetected= len(li)
    prev_grey, grey = grey, prev_grey
    prev_pyramid, pyramid = pyramid, prev_pyramid
    cv.ShowImage("Fig",segment)

    counter+=1 
    c = cv.WaitKey(10) % 0x100
    if c == 27: break 
    if 32 <= c and c < 128:
      cc = chr(c).lower()
      if cc== ' ':
        extract=True
      if cc=='r':
        extract=False
        selected=0
        colors=[[] for i in range(6)]
        didassignments=False
        assigned=[[-1 for i in range(9)] for j in range(6)]
        for i in range(6):
            assigned[i][4]=i    
        didassignments=False
        
      if cc=='n':
        selected=selected-1
        if selected<0: selected=5
      if cc=='m':
        selected=selected+1
        if selected>5: selected=0
      if cc=='d':
        dodetection=not dodetection
      if cc=='p':
        processColors()
      if cc=='s':
        cv.SaveImage("C:\\code\\img\\pic"+`time()`+".jpg",segmentc)
        
      
cv.DestroyWindow("Fig")
